# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: inkyu

from collections import namedtuple

from coin.strategy.hk_hard_arb.util import PassAggExchangeMap
from coin.strategy.mm.calculator import MovingSum

PassivePrices = namedtuple('PassivePrices',
                           ['askp', 'bidp', 'askp_pull', 'bidp_pull', 'edge_buy', 'edge_sell'])

AggressivePrices = namedtuple('AggressivePrices', ['true_book_price', 'posting_price', 'edge'])


class SingleProductPriceCalculator(object):
  def __init__(self, *, product_holder, book_askt_func, book_bidt_func):
    # TODO(inkyu): Remove tick from here
    self._tick = product_holder.get_tick()
    self._book_askt_func = book_askt_func
    self._book_bidt_func = book_bidt_func

    self.timestamp = None
    self.ready = False

    self.ask0 = 0.
    self.bid0 = 0.
    self.askt = 0.
    self.bidt = 0.
    self.midp = 0.
    self.midt = 0.
    self.ask0_minus_one = 0.
    self.bid0_plus_one = 0.

  def reset(self):
    self.ready = False

  def update(self, book):
    self.timestamp = book.timestamp

    if not (book.has_ask() and book.has_bid()):
      self.reset()
      return

    self.ask0 = book.ask0().price
    self.bid0 = book.bid0().price
    self.askt = self._book_askt_func(book)
    self.bidt = self._book_bidt_func(book)
    self.midp = (self.ask0 + self.bid0) / 2
    self.midt = (self.askt + self.bidt) / 2

    assert self.askt >= self.ask0 >= self.midp >= self.bid0 >= self.bidt, (self.askt,
                                                                           self.ask0,
                                                                           self.midp,
                                                                           self.bid0,
                                                                           self.bidt)

    # Cut-in-line
    self.ask0_minus_one = self._tick.get_prev_price(self.ask0)
    self.bid0_plus_one = self._tick.get_next_price(self.bid0)

    self.ready = True

  def correct_ask_passive(self, ask_orig, cut_in_line=False):
    askp = self._tick.ceil(ask_orig)
    if cut_in_line:
      askp = max(askp, self.ask0_minus_one, self.bid0_plus_one)
    return askp

  def correct_bid_passive(self, bid_orig, cut_in_line=False):
    bidp = self._tick.floor(bid_orig)
    if cut_in_line:
      bidp = min(bidp, self.ask0_minus_one, self.bid0_plus_one)
    return bidp

  def correct_ask_aggressive(self, ask_orig):
    return self._tick.get_prev_price(ask_orig)

  def correct_bid_aggressive(self, bid_orig):
    return self._tick.get_next_price(bid_orig)


class PriceCalculator(object):
  def __init__(self, *, basis_ma_window, product_holders, book_askt_funcs, book_bidt_funcs):
    self._use_basis = basis_ma_window > 0
    self._basis_ma_window = basis_ma_window
    self._basis_ma_required = 10

    self._prices_product = PassAggExchangeMap(
        passive=SingleProductPriceCalculator(product_holder=product_holders.passive,
                                             book_askt_func=book_askt_funcs.passive,
                                             book_bidt_func=book_bidt_funcs.passive),
        aggressive=SingleProductPriceCalculator(product_holder=product_holders.aggressive,
                                                book_askt_func=book_askt_funcs.aggressive,
                                                book_bidt_func=book_bidt_funcs.aggressive))

    self.ready = False

    if self._use_basis:
      self._ask_basis_ma = MovingSum(self._basis_ma_window)
      self._bid_basis_ma = MovingSum(self._basis_ma_window)

  def reset(self):
    self.ready = False

  def update(self, *, book_passive, book_aggressive):
    self._prices_product.passive.update(book_passive)
    self._prices_product.aggressive.update(book_aggressive)

    if not all(price_product.ready for price_product in self._prices_product):
      self.reset()
      return

    book_ts = max(price_product.timestamp for price_product in self._prices_product)
    old_ts = min(price_product.timestamp for price_product in self._prices_product)

    if book_ts - old_ts > 2.e9:  # TODO(inkyu): Parameterize
      self.reset()
      return

    if self._use_basis:
      askt_diff = self._prices_product.passive.askt - self._prices_product.aggressive.askt
      bidt_diff = self._prices_product.passive.bidt - self._prices_product.aggressive.bidt

      self._ask_basis_ma.update(book_ts, askt_diff)
      self._bid_basis_ma.update(book_ts, bidt_diff)

      if (len(self._ask_basis_ma) < self._basis_ma_required
          or len(self._bid_basis_ma) < self._basis_ma_required):
        self.reset()
        return

    self.ready = True

  def calculate_price_passive(self, edge, *, pull_edge=None, cut_in_line=True):
    assert self.ready
    assert -1 < edge < 1

    askt_agg_proj_pass = self._prices_product.aggressive.askt
    bidt_agg_proj_pass = self._prices_product.aggressive.bidt
    if self._use_basis:
      askt_agg_proj_pass += float(self._ask_basis_ma.average)
      bidt_agg_proj_pass += float(self._bid_basis_ma.average)

    askp_passive = self._prices_product.passive.correct_ask_passive(askt_agg_proj_pass *
                                                                    (1. + edge),
                                                                    cut_in_line=True)
    bidp_passive = self._prices_product.passive.correct_bid_passive(bidt_agg_proj_pass *
                                                                    (1. - edge),
                                                                    cut_in_line=True)

    if pull_edge is not None:
      askp_passive_pull = self._prices_product.passive.correct_ask_passive(askt_agg_proj_pass *
                                                                           (1. + pull_edge),
                                                                           cut_in_line=False)
      bidp_passive_pull = self._prices_product.passive.correct_bid_passive(bidt_agg_proj_pass *
                                                                           (1. - pull_edge),
                                                                           cut_in_line=False)
    else:
      askp_passive_pull = None
      bidp_passive_pull = None

    edge_buy = 1. - self._prices_product.passive.bidt / bidt_agg_proj_pass
    edge_sell = self._prices_product.passive.askt / askt_agg_proj_pass - 1.

    return PassivePrices(askp=askp_passive,
                         bidp=bidp_passive,
                         askp_pull=askp_passive_pull,
                         bidp_pull=bidp_passive_pull,
                         edge_buy=edge_buy,
                         edge_sell=edge_sell)

  def calculate_price_aggressive_buy(self, fill_price, agg_edge):  # fill_price = ask0_pass
    askt_agg = self._prices_product.aggressive.askt
    fill_price_proj_agg = fill_price
    if self._use_basis:
      fill_price_proj_agg -= float(self._ask_basis_ma.average)
    agg_buy_price = self._prices_product.aggressive.correct_bid_aggressive(fill_price_proj_agg *
                                                                           (1. - agg_edge))
    edge = fill_price_proj_agg / askt_agg - 1.

    return AggressivePrices(true_book_price=askt_agg, posting_price=agg_buy_price, edge=edge)

  def calculate_price_aggressive_sell(self, fill_price, agg_edge):  # fill_price = bid0_pass
    bidt_agg = self._prices_product.aggressive.bidt
    fill_price_proj_agg = fill_price
    if self._use_basis:
      fill_price_proj_agg -= float(self._bid_basis_ma.average)
    agg_sell_price = self._prices_product.aggressive.correct_ask_aggressive(fill_price_proj_agg *
                                                                            (1. + agg_edge))
    edge = 1. - fill_price_proj_agg / bidt_agg

    return AggressivePrices(true_book_price=bidt_agg, posting_price=agg_sell_price, edge=edge)
